package com.study.feature.extract

import org.apache.spark.ml.feature.{HashingTF, IDF, Tokenizer}
import org.apache.spark.sql.SparkSession

/**
 * 特征提取-TF-IDF(词频－逆向文档频率)
 *
 * @author stephen
 * @date 2019-08-28 09:36
 */
object TfIdfDemo {

  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName(this.getClass.getSimpleName)
      .master("local[*]")
      .getOrCreate()

    spark.sparkContext.setLogLevel("warn")

    val sentenceData = spark.createDataFrame(Seq(
      (0.0, "I heard about Spark and I love Spark"),
      (0.0, "I wish Java could use case classes"),
      (1.0, "Logistic regression models are neat")
    )).toDF("label", "sentence")

    // 定义分词器（也可以使用其他分词器）
    val tokenizer = new Tokenizer().setInputCol("sentence").setOutputCol("words")
    // 进行分词
    val wordsData = tokenizer.transform(sentenceData)

    // todo 去掉停用词

    // 定义TF，HashingTF 是一个Transformer
    val hashingTF = new HashingTF()
      .setInputCol("words")
      .setOutputCol("rawFeatures")
      // 设置哈希表的桶数(原始特征通过hash函数，映射到一个索引值。后面只需要统计这些索引值的频率，就可以知道对应词的频率。)
      .setNumFeatures(20)
    // HashingTF将句子转换为特征向量
    val featurizedData = hashingTF.transform(wordsData)

    // 定义IDF，IDF是一个Estimator
    val idf = new IDF().setInputCol("rawFeatures").setOutputCol("features")
    // 产生一个IDFModel
    val model = idf.fit(featurizedData)

    val rescaledData = model.transform(featurizedData)
    rescaledData.select("label", "words", "rawFeatures", "features").show(false)
    // 20：哈希表的桶数
    // [0,5,9,13,17]:词的哈希值(有一个单词是没有的)
    // [1.0,2.0,2.0,1.0,2.0]:词频
    // [0.6931471805599453,1.3862943611198906,0.5753641449035617,0.0,1.3862943611198906]: 对应单词的TF-IDF值

    spark.stop()
  }
}
